package cn.com.greatwall.kafka.util;

import java.io.Closeable;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;

import org.I0Itec.zkclient.ZkClient;
import org.I0Itec.zkclient.exception.ZkMarshallingError;
import org.I0Itec.zkclient.serialize.ZkSerializer;
import org.apache.zookeeper.data.Stat;

import com.alibaba.fastjson.JSONObject;

import cn.com.greatwall.kafka.service.dto.BrokerDto;
import kafka.utils.ZKStringSerializer;
import lombok.Data;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;

/**
 * @Author herw
 * @Time 2021-01-13 09:38:44
 * @Version 1.0
 * @Description: TODO(用一句话描述该文件做什么)
 */
@Slf4j
public class ZKUtil implements Closeable {
    private static final int SESSIONT_IMEOUT = 1000 * 60;
    private static final int CONNECTION_TIMEOUT = 1000 * 60;
    private ZkClient zkClient;
    private String rootPath = "";

    public ZKUtil(String zkServers) {
        String lastChar = zkServers.substring(zkServers.length() - 1, zkServers.length());
        // 如果最后一个字符是'/'，去掉
        if ("/".equals(lastChar)) {
            zkServers = zkServers.substring(0, zkServers.length() - 1);
        }
        // 获取指定path
        int index = zkServers.indexOf("/");
        if (index != -1) {
            // 存在rootPath
            String newzkServers = zkServers.substring(0, index);
            this.rootPath = zkServers.substring(index);
            zkClient = new ZkClient(newzkServers, SESSIONT_IMEOUT, CONNECTION_TIMEOUT, new ZkStringSerialize());
        } else {
            zkClient = new ZkClient(zkServers, SESSIONT_IMEOUT, CONNECTION_TIMEOUT, new ZkStringSerialize());
        }
    }
    
    public List<BrokerDto> getBrokers() throws Exception {
        List<BrokerDto> kafkaHosts = new ArrayList<BrokerDto>();
        try {
            List<String> ids = zkClient.getChildren(rootPath + "/brokers/ids");
            for (String id : ids) {
                String brokerInfo = readDataMaybeNull(rootPath + "/brokers/ids/" + id).getData();
                BrokerDto bi = new BrokerDto();
                bi.setBid(Integer.parseInt(id));
                JSONObject jsonObj = JSONObject.parseObject(brokerInfo);
                if (jsonObj.containsKey("host")) {
                    bi.setHost(jsonObj.get("host").toString());
                }
                if (jsonObj.containsKey("port")) {
                    bi.setPort(Integer.parseInt(jsonObj.get("port").toString()));
                }
                if (jsonObj.containsKey("jmx_port")) {
                    bi.setJmxPort(Integer.parseInt(jsonObj.get("jmx_port").toString()));
                }
                if (jsonObj.containsKey("version")) {
                    bi.setVersion(Integer.parseInt(jsonObj.get("version").toString()));
                }
                if (jsonObj.containsKey("timestamp")) {
                    bi.setTimestamp(Long.parseLong(jsonObj.get("timestamp").toString()));
                }
                kafkaHosts.add(bi);
            }
        } catch (Exception e) {
            // TODO: handle exception
        }
        return kafkaHosts;
    }
    
    private ZkDataAndStat readDataMaybeNull(String path) {
        Stat stat = new Stat();
        String data = null;
        try {
            data = zkClient.readData(path, stat);
        } catch (Exception e) {
            log.warn("zk read data error." ,e);
        }
        return new ZkDataAndStat(data, stat);
    }
    
    public JSONObject descriptionConfig(String topicName) {
        String configInfo = this.readDataMaybeNull(rootPath + "/config/topics/" + topicName).getData();
        JSONObject configObject = JSONObject.parseObject(configInfo);
        
        return configObject.getJSONObject("config");
    }

    @Override
    public void close() throws IOException {
        if (zkClient != null) {
            zkClient.close();
        }
    }

    class ZkStringSerialize implements ZkSerializer {

        @Override
        public byte[] serialize(Object data) throws ZkMarshallingError {
            return ZKStringSerializer.serialize(data);
        }

        @Override
        public Object deserialize(byte[] bytes) throws ZkMarshallingError {
            return ZKStringSerializer.deserialize(bytes);
        }
    }
    
    @Data
    class ZkDataAndStat {
        private String data;
        private Stat stat;
        
        public ZkDataAndStat(String data, Stat stat) {
            super();
            this.data = data;
            this.stat = stat;
        }
    }
}
